#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@文件名:   librosaLoad.py    
@作者:     haoge 
@时间:   2021/4/7 15:30

@函数名             @版本         @功能描述
------------      ---------    --------    



'''
import os
import random

import pandas as pd
import librosa
import tensorflow as tf
import time
from tqdm import tqdm
import numpy as np
class_dim = 20
EPOCHS = 100
BATCH_SIZE=32
MyGeneralPath= "C:\\Users\\haoge\\Desktop\\BirdsData"#数据集总路径
MyLabelDir= "C:\\Users\\haoge\\Desktop\\BirdsLabelList.txt"
MyBirdsDataDir="C:\\Users\\haoge\\Desktop\\MyBirdsData"
myLabelList={}
myLabelList["0009"]= "灰雁"
myLabelList["0017"]= "大天鹅"
myLabelList["0034"]= "绿头鸭"
myLabelList["0036"]= "绿翅鸭"
myLabelList["0074"]= "灰山鹑"
myLabelList["0077"]= "西鹌鹑"
myLabelList["0114"]= "雉鸡"
myLabelList["0121"]= "红喉潜鸟"
myLabelList["0180"]= "苍鹭"
myLabelList["0202"]= "普通鸬鹚"
myLabelList["0235"]= "苍鹰"
myLabelList["0257"]= "欧亚鵟"
myLabelList["0265"]= "西方秧鸡"
myLabelList["0281"]= "骨顶鸡"
myLabelList["0298"]= "黑翅长脚鹬"
myLabelList["0300"]= "凤头麦鸡"
myLabelList["0364"]= "白腰草鹬"
myLabelList["0368"]= "红脚鹬"
myLabelList["0370"]= "林鹬"
myLabelList["1331"]= "麻雀"
class BirdsDirStruct:
    def __init__(self,labelList,kindsList,kindsMap):
        self.labelList=labelList
        self.kindsList=kindsList
        self.kindsMap=kindsMap

def GetBirdsDir(GeneralPath):

    kindsList=os.listdir(GeneralPath)


    kindsMap={}
    for i in kindsList:
        kindList=os.listdir(GeneralPath+"\\"+i)
        oneWav=[]
        for j in kindList:
            oneWav.append(GeneralPath+"\\"+i+"\\"+j)
            # oneWav.append(j[7])
        kindsMap[i]=oneWav
    BirdsDirStructData=BirdsDirStruct(myLabelList, kindsList, kindsMap)
    return BirdsDirStructData
BirdsDirData_Struct = GetBirdsDir(MyGeneralPath)
kindsList=list(BirdsDirData_Struct.labelList.keys())#009,0281....

#==========================================================================================
def get_data_list(audio_path, list_path):
    sound_sum = 0
    audios = os.listdir(audio_path)
    f_train = open(os.path.join(list_path, 'train_list.txt'), 'w')
    f_test = open(os.path.join(list_path, 'test_list.txt'), 'w')

    for i in range(len(audios)):
        sounds = os.listdir(os.path.join(audio_path, audios[i]))
        for sound in sounds:
            sound_path = os.path.join(audio_path, audios[i], sound)
            t = librosa.get_duration(filename=sound_path)
            # [可能需要修改参数] 过滤小于2.1秒的音频
            if t ==2:
                if sound_sum % 100 == 0:
                    f_test.write('%s\t%d\n' % (sound_path, i))
                else:
                    f_train.write('%s\t%d\n' % (sound_path, i))
                sound_sum += 1
        print("Audio：%d/%d" % (i + 1, len(audios)))

    f_test.close()
    f_train.close()


# get_data_list('C:/Users/haoge/Desktop/dataset/audio', 'C:/Users/haoge/Desktop/dataset')

# 获取浮点数组
def _float_feature(value):
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))


# 获取整型数据
def _int64_feature(value):
    if not isinstance(value, list):
        value = [value]
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))


# 把数据添加到TFRecord中
def data_example(data, label):
    feature = {
        'data': _float_feature(data),
        'label': _int64_feature(label),
    }
    return tf.train.Example(features=tf.train.Features(feature=feature))


# 开始创建tfrecord数据
def create_data_tfrecord(data_list_path, save_path):
    with open(data_list_path, 'r') as f:
        data = f.readlines()
    with tf.io.TFRecordWriter(save_path) as writer:
        for d in tqdm(data):
            try:
                path, label = d.replace('\n', '').split('\t')
                wav, sr = librosa.load(path, sr=16000)
                intervals = librosa.effects.split(wav, top_db=20)
                wav_output = []
                # [可能需要修改参数] 音频长度 16000 * 秒数
                wav_len = int(16000 * 2)
                for sliced in intervals:
                    wav_output.extend(wav[sliced[0]:sliced[1]])
                for i in range(5):
                    # 裁剪过长的音频，过短的补0
                    # if len(wav_output) > wav_len:
                    #     l = len(wav_output) - wav_len
                    #     r = random.randint(0, l)
                    #     wav_output = wav_output[r:wav_len + r]
                    # else:
                    #     wav_output.extend(np.zeros(shape=[wav_len - len(wav_output)], dtype=np.float32))
                    wav_output = np.array(wav_output)
                    # 转成梅尔频谱
                    ps = librosa.feature.melspectrogram(y=wav_output, sr=sr, hop_length=256).reshape(-1).tolist()
                    # [可能需要修改参数] 梅尔频谱shape ，librosa.feature.melspectrogram(y=wav_output, sr=sr, hop_length=256).shape
                    if len(ps) != 15360: continue
                    tf_example = data_example(ps, int(label))
                    writer.write(tf_example.SerializeToString())
                    if len(wav_output) <= wav_len:
                        break
            except Exception as e:
                print(e)



create_data_tfrecord("C:\\Users\\haoge\\Desktop\\dataset\\train_list.txt", "C:\\Users\\haoge\\Desktop\\dataset\\train.tfrecord")
create_data_tfrecord("C:\\Users\\haoge\\Desktop\\dataset\\test_list.txt", "C:\\Users\\haoge\\Desktop\\dataset\\test.tfrecord")






















# for i in BirdsDirData_Struct.labelList.keys():
#     print("i:===========================================", i)
#     print("========================================", len(BirdsDirData_Struct.kindsMap[i]))
#     for j in range(len(BirdsDirData_Struct.kindsMap[i])):
#         # ================加载音频文件============================================
#         y1, sr1 = librosa.load(BirdsDirData_Struct.kindsMap[i][j])  # oneWav就是每一个类别音频WAV的数据，sr就是采样率22050，加载音频文件
